// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/lib/array_shuffle.h>
#include <pollux/functions/lib/registration_helpers.h>
#include <pollux/functions/lib/repeat.h>
#include <pollux/functions/lib/slice.h>
#include <pollux/functions/prestosql/array_functions.h>
#include <pollux/functions/sparksql/array_append.h>
#include <pollux/functions/sparksql/array_concat.h>
#include <pollux/functions/sparksql/array_flatten_function.h>
#include <pollux/functions/sparksql/array_insert.h>
#include <pollux/functions/sparksql/array_min_max_function.h>
#include <pollux/functions/sparksql/array_sort.h>

namespace kumo::pollux::functions {
    // POLLUX_REGISTER_VECTOR_FUNCTION must be invoked in the same namespace as the
    // vector function definition.
    // Higher order functions.
    void registerSparkArrayFunctions(const std::string &prefix) {
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_transform, prefix + "transform");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_reduce, prefix + "aggregate");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_constructor, prefix + "array");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_contains, prefix + "array_contains");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_distinct, prefix + "array_distinct");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_except, prefix + "array_except");
        POLLUX_REGISTER_VECTOR_FUNCTION(
            udf_array_intersect, prefix + "array_intersect");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_position, prefix + "array_position");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_zip, prefix + "arrays_zip");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_any_match, prefix + "exists");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_filter, prefix + "filter");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_all_match, prefix + "forall");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_zip_with, prefix + "zip_with");
    }

    namespace sparksql {
        template<typename T>
        inline void registerArrayConcatFunction(const std::string &prefix) {
            register_function<
                ParameterBinder<ArrayConcatFunction, T>,
                Array<T>,
                Variadic<Array<T> > >({prefix + "concat"});
        }

        void registerArrayConcatFunctions(const std::string &prefix) {
            registerArrayConcatFunction<int8_t>(prefix);
            registerArrayConcatFunction<int16_t>(prefix);
            registerArrayConcatFunction<int32_t>(prefix);
            registerArrayConcatFunction<int64_t>(prefix);
            registerArrayConcatFunction<int128_t>(prefix);
            registerArrayConcatFunction<float>(prefix);
            registerArrayConcatFunction<double>(prefix);
            registerArrayConcatFunction<bool>(prefix);
            registerArrayConcatFunction<Varbinary>(prefix);
            registerArrayConcatFunction<Varchar>(prefix);
            registerArrayConcatFunction<Timestamp>(prefix);
            registerArrayConcatFunction<Date>(prefix);
            registerArrayConcatFunction<UnknownValue>(prefix);
            registerArrayConcatFunction<Generic<T1> >(prefix);
        }

        inline void registerArrayJoinFunctions(const std::string &prefix) {
            register_function<
                ParameterBinder<ArrayJoinFunction, Varchar>,
                Varchar,
                Array<Varchar>,
                Varchar>({prefix + "array_join"});

            register_function<
                ParameterBinder<ArrayJoinFunction, Varchar>,
                Varchar,
                Array<Varchar>,
                Varchar,
                Varchar>({prefix + "array_join"});
        }

        template<typename T>
        inline void registerArrayMinMaxFunctions(const std::string &prefix) {
            register_function<ArrayMinFunction, T, Array<T> >({prefix + "array_min"});
            register_function<ArrayMaxFunction, T, Array<T> >({prefix + "array_max"});
        }

        inline void registerArrayMinMaxFunctions(const std::string &prefix) {
            registerArrayMinMaxFunctions<int8_t>(prefix);
            registerArrayMinMaxFunctions<int16_t>(prefix);
            registerArrayMinMaxFunctions<int32_t>(prefix);
            registerArrayMinMaxFunctions<int64_t>(prefix);
            registerArrayMinMaxFunctions<int128_t>(prefix);
            registerArrayMinMaxFunctions<float>(prefix);
            registerArrayMinMaxFunctions<double>(prefix);
            registerArrayMinMaxFunctions<bool>(prefix);
            registerArrayMinMaxFunctions<Varbinary>(prefix);
            registerArrayMinMaxFunctions<Varchar>(prefix);
            registerArrayMinMaxFunctions<Timestamp>(prefix);
            registerArrayMinMaxFunctions<Date>(prefix);
            registerArrayMinMaxFunctions<Orderable<T1> >(prefix);
        }

        template<typename T>
        inline void registerArrayRemoveFunctions(const std::string &prefix) {
            register_function<ArrayRemoveFunction, Array<T>, Array<T>, T>(
                {prefix + "array_remove"});
        }

        inline void registerArrayRemoveFunctions(const std::string &prefix) {
            registerArrayRemoveFunctions<int8_t>(prefix);
            registerArrayRemoveFunctions<int16_t>(prefix);
            registerArrayRemoveFunctions<int32_t>(prefix);
            registerArrayRemoveFunctions<int64_t>(prefix);
            registerArrayRemoveFunctions<int128_t>(prefix);
            registerArrayRemoveFunctions<float>(prefix);
            registerArrayRemoveFunctions<double>(prefix);
            registerArrayRemoveFunctions<bool>(prefix);
            registerArrayRemoveFunctions<Timestamp>(prefix);
            registerArrayRemoveFunctions<Date>(prefix);
            registerArrayRemoveFunctions<Varbinary>(prefix);
            registerArrayRemoveFunctions<Generic<T1> >(prefix);
            register_function<
                ArrayRemoveFunctionString,
                Array<Varchar>,
                Array<Varchar>,
                Varchar>({prefix + "array_remove"});
        }

        template<typename T>
        inline void registerArrayUnionFunction(const std::string &prefix) {
            register_function<ArrayUnionFunction, Array<T>, Array<T>, Array<T> >(
                {prefix + "array_union"});
        }

        inline void registerArrayUnionFunctions(const std::string &prefix) {
            registerArrayUnionFunction<int8_t>(prefix);
            registerArrayUnionFunction<int16_t>(prefix);
            registerArrayUnionFunction<int32_t>(prefix);
            registerArrayUnionFunction<int64_t>(prefix);
            registerArrayUnionFunction<int128_t>(prefix);
            registerArrayUnionFunction<float>(prefix);
            registerArrayUnionFunction<double>(prefix);
            registerArrayUnionFunction<bool>(prefix);
            registerArrayUnionFunction<Timestamp>(prefix);
            registerArrayUnionFunction<Date>(prefix);
            registerArrayUnionFunction<Varbinary>(prefix);
            registerArrayUnionFunction<Varchar>(prefix);
            registerArrayUnionFunction<Generic<T1> >(prefix);
        }

        void registerArrayFunctions(const std::string &prefix) {
            registerArrayConcatFunctions(prefix);
            registerArrayJoinFunctions(prefix);
            registerArrayMinMaxFunctions(prefix);
            registerArrayRemoveFunctions(prefix);
            registerSparkArrayFunctions(prefix);
            // Register array sort functions.
            exec::registerStatefulVectorFunction(
                prefix + "array_sort", arraySortSignatures(), makeArraySort);
            exec::registerStatefulVectorFunction(
                prefix + "sort_array", sortArraySignatures(), makeSortArray);
            exec::registerStatefulVectorFunction(
                prefix + "array_repeat",
                repeatSignatures(),
                makeRepeatAllowNegativeCount,
                repeatMetadata());
            register_function<
                ArrayFlattenFunction,
                Array<Generic<T1> >,
                Array<Array<Generic<T1> > > >({prefix + "flatten"});
            register_function<
                ArrayInsert,
                Array<Generic<T1> >,
                Array<Generic<T1> >,
                int32_t,
                Generic<T1>,
                bool>({prefix + "array_insert"});
            POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_get, prefix + "get");
            exec::registerStatefulVectorFunction(
                prefix + "shuffle",
                arrayShuffleWithCustomSeedSignatures(),
                makeArrayShuffleWithCustomSeed,
                getMetadataForArrayShuffle());
            registerIntegerSliceFunction(prefix);
            register_function<
                ArrayAppendFunction,
                Array<Generic<T1> >,
                Array<Generic<T1> >,
                Generic<T1> >({prefix + "array_append"});
            registerArrayUnionFunctions(prefix);
        }
    } // namespace sparksql
} // namespace kumo::pollux::functions
